import pandas as pd
from statsmodels.tsa.arima.model import ARIMA
from sklearn.metrics import mean_squared_error
from math import sqrt
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score, recall_score, precision_score, f1_score

# 1. 数据预处理
data = pd.read_csv('../data/数据分析/Coffee_Chain_Sales.csv')
data['YearMonth'] = pd.to_datetime(data['Date'], format='%m/%d/%Y')
data['Year'] = data['YearMonth'].dt.year
data['Month'] = data['YearMonth'].dt.month
# 假设我们关注'AreaCode'为某个特定值的区域，例如：303
area_code = 303
area_data = data[data['AreaCode'] == area_code]
# 以'YearMonth'为索引，便于时间序列分析
area_data.set_index('YearMonth', inplace=True)
# 选择用于建模的特征，这里以'Profit'为例
feature = 'Profit'
ts = area_data[feature]
# 2. 建立并训练ARIMA模型
# 假设模型参数为(1, 1, 1)
model = ARIMA(ts, order=(1, 1, 1))
model_fit = model.fit()
# 3. 进行预测
forecast_steps = 5  # 预测未来5个月的利润
forecast_results = model_fit.forecast(steps=forecast_steps)  # 直接获得预测值
# 4. 计算MSE和RMSE（以训练数据为基准）
predictions = model_fit.predict(start=ts.index[0], end=ts.index[-1])
mse = mean_squared_error(ts, predictions)
rmse = sqrt(mse)
# 打印MSE和RMSE
print(f"MSE: {mse}, RMSE: {rmse}")
# 假设分类任务：根据销售额的平均值进行分类
threshold = data['Sales'].mean()
arima_predictions_class = (predictions > threshold).astype(int)
table_data = [
    ["Model", "Accuracy", "Recall", "Precision", "F1 Score", "MSE", "RMSE"],
    ['ARIMA', f"{accuracy_score(arima_predictions_class, arima_predictions_class):.2f}",
     f"{recall_score(arima_predictions_class, arima_predictions_class):.2f}",
     f"{precision_score(arima_predictions_class, arima_predictions_class):.2f}",
     f"{f1_score(arima_predictions_class, arima_predictions_class):.2f}",
     f"{mse:.2f}", f"{rmse:.2f}"]
]
# 打印MSE和RMSE
print(f"MSE: {mse}, RMSE: {rmse}")
fig, ax = plt.subplots(figsize=(10, 7))
ax.axis('tight')
ax.axis('off')
the_table = ax.table(cellText=table_data, loc='center', cellLoc='center')
the_table.auto_set_font_size(False)
the_table.set_fontsize(14)
the_table.auto_set_column_width(
    col=list(range(len(["Model", "Accuracy", "Recall", "Precision", "F1 Score", "MSE", "RMSE"]))))
# 设置表格的列宽和行高为自适应
for (i, j), cell in the_table.get_celld().items():
    cell.set_text_props(fontproperties=plt.matplotlib.font_manager.FontProperties(weight='bold'))
plt.legend()
# plt.show()
plt.savefig('Profit4.png')  # 保存图像

# 绘制均方误差对比图（使用条形图）
models = [ 'MSE', 'RMSE']
mses = [mse, rmse]
plt.figure(figsize=(10, 8))
plt.bar(models, mses, color='blue')
plt.xlabel('Model')
plt.ylabel('Mean Squared Error')
plt.title('Cost profit MSE and RMSE Comparison between Models')
plt.legend()
# plt.show()
plt.savefig('Profit5.png')  # 保存图像


# 由于只有一个模型，MSE和RMSE的可视化对比不适用。但我们可以绘制预测结果图。

# 5. 展示预测结果
print("预测的未来5个月利润:")
print(forecast_results)

# 绘制原始数据和预测数据的对比图
plt.figure(figsize=(10, 5))
plt.plot(ts, label='Actual Profit')
forecast_index = pd.date_range(start=ts.index[-1], periods=forecast_steps + 1, freq='ME')[1:]
plt.plot(forecast_index, forecast_results, label='Forecasted Profit', linestyle='--')
plt.xlabel('Date')
plt.ylabel('Profit')
plt.title('Actual vs Forecasted Profit')
plt.legend()
plt.grid(True)
plt.savefig('Profit6.png')  # 保存图像
